﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace Oops.Xna.Framework
{
	/// <summary>
	/// An extension to the <see cref="System.IEnumerable<T>" /> interface.
	/// </summary>
	public static class EnumerableExtension
	{
		public static IEnumerable<TSource> SelectRecursive<TSource>(
			this IEnumerable<TSource> source, 
			Func<TSource, IEnumerable<TSource>> recursiveSelector)
		{
			// TODO : GC!
			// Create a stack of enumerators and add the source's as the first item.
			Stack<IEnumerator<TSource>> stack = new Stack<IEnumerator<TSource>>();
			stack.Push(source.GetEnumerator());

			try
			{
				// Iterate until there is no more items on the stack.
				while (stack.Count > 0)
				{
					// Item extis.
					if (stack.Peek().MoveNext())
					{
						// Get the the item from the current enumerator and yield it as the next item.
						TSource current = stack.Peek().Current;
						yield return current;

						// Add the current item's enumerator onto the stack.
						stack.Push(recursiveSelector(current).GetEnumerator());
					}
					else
					{
						// Dispose the current enumerator on the stack and remove it.
						stack.Pop().Dispose();
					}
				}
			}
			finally
			{
				// Dispose of all remaining enumerators on the stack.
				while (stack.Count > 0)
				{
					stack.Pop().Dispose();
				}
			}
		}
	}
}
